package main.state

import main.models.{AbstractEntity, StateAction, World}
import main.utils.Tree

// Class representing world state
// TODO: Combine State and lens. Add ability to store state next to stateLens to avoid boilerplate

case class State(
    world: StateLens[World],
    ownershipTree: Tree[StateLens[? <: AbstractEntity]],
    entities: Map[StateLens[AbstractEntity], AbstractEntity] = Map()
) {

  def update(stateLens: StateLens[AbstractEntity], newEntity: AbstractEntity): State =
    copy(entities = entities.updated(stateLens, newEntity))

  def registerEntity[T <: AbstractEntity](entity: T): (StateLens[T], State) = {
    val newLens = new StateLens[T] {
      val id: Int = entity.id
      override def get(implicit state: State): T = state.entities(this).asInstanceOf[T]
    }
    (newLens, copy(entities = entities + (newLens -> entity)))
  }


  def apply(stateAction: StateAction): State = {
    stateAction(this)
  }

}

object State {
  def fromWorld(world: World): (StateLens[World], State) = {
    val worldLens = new StateLens[World] {
      val id: Int = world.id
      override def get(implicit state: State): World = state.entities(this).asInstanceOf[World]
    }
    (worldLens, State(worldLens, new Tree(worldLens)))
  }
}

// Interface representing entity lens
// TODO: Fix reference checking

trait StateLens[+T <: AbstractEntity] {

  val id: Int

  def get(implicit state: State): T

  def map[U >: T <: AbstractEntity](implicit state: State, fn: U => U): State = {
    state.update(this, fn(this.get.asInstanceOf[U]))
  }

  def addChild[U <: AbstractEntity](implicit state: State, entityToAdd: StateLens[U]): State =
    state.copy(ownershipTree = state.ownershipTree.addChild(this, entityToAdd))

  def removeChild[U <: AbstractEntity](implicit state: State, entityToRemove: StateLens[U]): State =
    state.copy(ownershipTree = state.ownershipTree.removeChild(this, entityToRemove))

  override def equals(obj: Any): Boolean = {
    if(!obj.isInstanceOf[StateLens[T]]) return false
    obj.asInstanceOf[StateLens[T]].id == id
  }

}
